package cn.seecoder.ai.algorithm.ml.regression

import cn.seecoder.ai.algorithm.ml.BaseRegression
import cn.seecoder.ai.enums.TrainParamEnum
import cn.seecoder.ai.model.bo.TrainParamsBO
import cn.seecoder.ai.utils.TrainParamReader.readParam
import org.apache.spark.ml.PipelineStage
import org.apache.spark.ml.regression.DecisionTreeRegressor
import org.slf4j.Logger

import java.util

class DecisionTreeRegression extends BaseRegression{


  override def buildMachineLearningStage(trainParams: TrainParamsBO): PipelineStage = {


    // Train a DecisionTree model.
    val dt = new DecisionTreeRegressor()
      .setLabelCol("label")
      .setFeaturesCol("indexedFeatures")
      .setSeed(trainParams.seed)
      .setMaxDepth(trainParams.maxDepth)
      .setMaxBins(trainParams.maxBins)
      .setMinInfoGain(trainParams.minInfoGain)

    dt
  }
}
